import torch
import torch.nn as nn
import torch.optim as optim
from PIL import Image
from torchvision import datasets, transforms

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 1. 加载数据
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)


# 2. 定义模型
class SoftmaxRegression(nn.Module):
    def __init__(self):
        super(SoftmaxRegression, self).__init__()
        self.linear = nn.Linear(28 * 28, 10)  # 输入是展平的784个像素，输出是10个类别

    def forward(self, x):
        x = x.view(-1, 28 * 28)  # 展平输入
        out = self.linear(x)
        return out


model = SoftmaxRegression().to(device)

# 3. 定义损失函数和优化器
loss_fn = nn.CrossEntropyLoss()  # CrossEntropyLoss 包含了 Softmax
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 4. 训练模型
num_epochs = 5
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images, labels = images.to(device), labels.to(device)
        # 前向传播
        outputs = model(images)
        loss = loss_fn(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], Loss: {loss.item():.4f}')

# 输出示例：
# Epoch [1/5], Step [100/938], Loss: 0.5163
# Epoch [1/5], Step [200/938], Loss: 0.3648
# ...
# 定义与训练时相同的图像转换操作
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),  # 转换为灰度图像
    transforms.ToTensor(),                        # 转换为张量
    transforms.Normalize((0.1307,), (0.3081,))    # 标准化，使用与训练集相同的均值和标准差
])

# 2. 加载并预处理单张图片
image_path = './data/figure.png'  # 替换为你的图片路径
image = Image.open(image_path)
image = transform(image).unsqueeze(0).to(device)  # 添加批量维度 (1, 1, 28, 28)

# 3. 使用模型进行预测
model.eval()  # 设置模型为评估模式
with torch.no_grad():
    output = model(image)  # 前向传播计算输出
    prediction = torch.argmax(output, dim=1)  # 获取最大值对应的类别

# 4. 输出预测结果
print(f'Predicted class: {prediction.item()}')